#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : mrna_get_gene_name.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/11/18 9:31
@Description :
    从 GTF 注释文件中提取 gene_id 与 gene_name 的对应关系，
    并将其补充到差异分析结果（TSV 文件）中。
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
from typing import Dict
import re


# ------------------------------------------------------------
# 命令行参数解析
# ------------------------------------------------------------
def _parse_args() -> Namespace:
    """
    解析命令行参数。

    Returns
    -------
    Namespace
        包含 deg_file、gtf_file、out_file 参数信息。
    """
    parser = ArgumentParser(description="根据 gene_id 从 GTF 文件获取 gene_name")

    parser.add_argument(
        "--deg_file",
        type=str,
        required=True,
        help="差异分析结果文件（TSV 格式，有表头，首列必须为 gene_id）"
    )

    parser.add_argument(
        "--gtf_file",
        type=str,
        required=True,
        help="基因组注释文件（GTF 格式）"
    )

    parser.add_argument(
        "--out_file",
        type=str,
        required=True,
        help="输出结果文件（TSV 格式，有表头，将自动创建目录）"
    )

    return parser.parse_args()


# ------------------------------------------------------------
# 解析 GTF 文件，生成 gene_id -> gene_name 映射
# ------------------------------------------------------------
def _parse_gtf(gtf_file: Path) -> Dict[str, str]:
    """
    解析 GTF 文件，提取 gene_id 与 gene_name 对应关系。

    Parameters
    ----------
    gtf_file : Path
        GTF 格式注释文件路径。

    Returns
    -------
    Dict[str, str]
        一个以 gene_id 为键、gene_name 为值的字典。

    Notes
    -----
    仅处理 feature 字段为 "gene" 的行；
    使用正则在注释字段（第 9 列）提取 gene_id 与 gene_name。
    """
    out_dict = {}

    with open(gtf_file, "r", encoding="utf-8") as handle:
        for line in handle:
            if line.startswith("#"):
                continue

            if line := line.strip():
                items = line.split("\t")
                anno = items[-1]

                # 仅解析类型为 gene 的记录
                if items[2] == "gene":
                    match = re.search(
                        r"gene_id\s+\"(.*?)\";.*?gene_name\s+\"(.*?)\";",
                        anno
                    )
                    if match:
                        gene_id = match.group(1)
                        gene_name = match.group(2)
                        out_dict[gene_id] = gene_name

    return out_dict


# ------------------------------------------------------------
# 根据 gene_id 注释 gene_name 并输出结果文件
# ------------------------------------------------------------
def _annotate_gene_id(deg_file: Path, anno_dict: Dict[str, str], out_file: Path) -> None:
    """
    为差异分析结果文件中的 gene_id 添加 gene_name 注释。

    Parameters
    ----------
    deg_file : Path
        输入差异表达结果文件（TSV）。
    anno_dict : Dict[str, str]
        gene_id 到 gene_name 的映射字典。
    out_file : Path
        输出的 TSV 文件路径。

    Notes
    -----
    输出文件表头固定为：
        gene_id  gene_name  baseMean  log2FoldChange  lfcSE  stat  pvalue  padj
    """
    with open(deg_file, "r", encoding="utf-8") as handle_r:
        with open(out_file, "w", encoding="utf-8") as handle_w:

            # 写入标准输出表头
            handle_w.write(
                "gene_name\tgene_id\tbaseMean\tlog2FoldChange\tlfcSE\tstat\tpvalue\tpadj\n"
            )

            # 跳过输入文件表头
            for line in handle_r.readlines()[1:]:
                if line := line.strip():
                    items = line.split("\t")
                    gene_id = items[0]
                    gene_name = anno_dict.get(gene_id, "--")  # 保持逻辑不变：缺失时会 KeyError

                    # 输出：gene_name + 原始整行内容
                    handle_w.write(f"{gene_name}\t{line}\n")


# ------------------------------------------------------------
# 主函数
# ------------------------------------------------------------
def main(args: Namespace) -> None:
    """
    主函数入口。
    负责解析路径、构建输出目录及生成 gene_id->gene_name 映射。
    """
    deg_file = Path(args.deg_file).absolute()
    gtf_file = Path(args.gtf_file).absolute()
    out_file = Path(args.out_file).absolute()

    # 确保输出目录存在
    out_file.parent.mkdir(parents=True, exist_ok=True)

    # 解析 GTF 文件
    out_dict = _parse_gtf(gtf_file)

    _annotate_gene_id(deg_file, out_dict, out_file)


if __name__ == "__main__":
    main(_parse_args())
